package main

import (
    "bytes"
    "crypto/des"
    "encoding/base64"
    "encoding/json"
    "flag"
    "fmt"
    "io"
    "io/ioutil"
    "log"
    "net/http"
    "os"
    "strings"
    "time"
)

const defaultDesKey = "ourheyday-yuandaima"

// Callback JSON structure from supervisor-agent/dispatch-agent
type CallbackPayload struct {
    TaskID    string `json:"task_id"`
    DeviceSN  string `json:"device_sn"`
    Stdout    string `json:"stdout"`
    Stderr    string `json:"stderr"`
    ExitCode  int    `json:"exit_code"`
    Timeout   bool   `json:"timeout"`
}

// API3 encrypted data structure
type Api3Data struct {
    Status  int    `json:"status"`
    ID      string `json:"id"`
    OhdCode string `json:"ohdCode"`
    Remark  string `json:"remark"`
}

func main() {
    listen := flag.String("listen", ":8008", "HTTP listen address (default :8008)")
    apiURL := flag.String("api-url", "", "API 3 endpoint URL (required)")
    desKey := flag.String("des-key", defaultDesKey, "DES encryption key, first 8 bytes used")
    debug := flag.Bool("debug", false, "Enable debug mode with full logging")
    flag.Parse()

    if *apiURL == "" {
        log.Println("ERROR: --api-url is required")
        flag.Usage()
        os.Exit(2)
    }
    if len(*desKey) < 8 {
        log.Println("ERROR: --des-key must be at least 8 bytes")
        os.Exit(2)
    }
    log.Printf("[*] callback-adapter listening at %s, forwarding to: %s", *listen, *apiURL)
    if *debug {
        log.Printf("[*] Debug mode enabled")
    }

    http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        start := time.Now()

        if r.Method != "POST" {
            http.Error(w, "Only POST allowed", http.StatusMethodNotAllowed)
            return
        }

        // Read request body
        bodyBytes, err := ioutil.ReadAll(r.Body)
        if err != nil {
            log.Printf("ERROR: Failed reading request body: %v", err)
            http.Error(w, err.Error(), http.StatusBadRequest)
            return
        }
        r.Body = ioutil.NopCloser(bytes.NewReader(bodyBytes))

        // Log - always print incoming
        if *debug {
            log.Println("===== Incoming Callback Start =====")
            log.Printf("Request URI: %s", r.RequestURI)
            log.Printf("Request headers:")
            for k, vs := range r.Header {
                for _, v := range vs {
                    log.Printf("  %s: %s", k, v)
                }
            }
            log.Printf("Request body: %s", string(bodyBytes))
            log.Println("===== Incoming Callback End =====")
        } else {
            log.Printf("[RECEIVED] URI=%s Content-Length=%d", r.RequestURI, len(bodyBytes))
        }

        // Parse incoming callback payload
        var cb CallbackPayload
        if err := json.Unmarshal(bodyBytes, &cb); err != nil {
            log.Printf("ERROR: Invalid callback JSON: %v, body: %s", err, string(bodyBytes))
            http.Error(w, err.Error(), http.StatusBadRequest)
            return
        }

        // Extract the "target_addr" field, prefer header, else URL param, then remove port
        targetAddr := r.Header.Get("target_addr")
        if targetAddr == "" {
            targetAddr = r.URL.Query().Get("target_addr")
        }
        ohdCode := targetAddr
        if idx := strings.Index(targetAddr, ":"); idx > 0 {
            ohdCode = targetAddr[:idx]
        }
        if ohdCode == "" {
            ohdCode = cb.DeviceSN
        }

        // Compute status: 1=success, 0=failure
        status := 1
        if cb.ExitCode != 0 || cb.Timeout {
            status = 0
        }

        // Assemble remark JSON: headers + raw payload
        headers := make(map[string][]string)
        for k, v := range r.Header {
            headers[k] = v
        }
        remarkObj := map[string]interface{}{
            "headers": headers,
            "body":    json.RawMessage(bodyBytes),
        }
        remarkJSONBytes, _ := json.Marshal(remarkObj)

        // Compose API3 data struct to encrypt
        api3 := Api3Data{
            Status:  status,
            ID:      cb.TaskID,
            OhdCode: ohdCode,
            Remark:  string(remarkJSONBytes),
        }
        api3Bytes, _ := json.Marshal(api3)

        if *debug {
            log.Println("===== API 3 Data to Encrypt =====")
            log.Printf("%s", string(api3Bytes))
            log.Println("=================================")
        }

        encryptedData, err := desEncrypt(api3Bytes, (*desKey)[:8])
        if err != nil {
            log.Printf("ERROR: DES encryption failed: %v", err)
            writeApi3Error(w, 500, "DES encryption failed: "+err.Error())
            return
        }

        if *debug {
            decryptedBack, _ := desDecrypt(encryptedData, (*desKey)[:8])
            log.Printf("Encrypted (base64): %s", encryptedData)
            log.Printf("Decrypted (again):  %s", decryptedBack)
        }

        // Compose outbound request to API3
        reqBody := map[string]interface{}{
            "result": true,
            "data":   encryptedData,
        }
        reqBodyBytes, _ := json.Marshal(reqBody)
        if *debug {
            log.Println("===== Outbound API 3 Request Start =====")
            log.Printf("POST %s", *apiURL)
            log.Printf("Content-Type: application/json")
            log.Printf("Body: %s", string(reqBodyBytes))
            log.Println("===== Outbound API 3 Request End =====")
        }

        reqAPI3, err := http.NewRequest("POST", *apiURL, bytes.NewReader(reqBodyBytes))
        reqAPI3.Header.Set("Content-Type", "application/json")
        httpClient := &http.Client{Timeout: 12 * time.Second}
        resp, err := httpClient.Do(reqAPI3)
        if err != nil {
            log.Printf("ERROR: Failed forwarding to API 3: %v", err)
            writeApi3Error(w, 502, "Forwarding error: "+err.Error())
            return
        }
        defer resp.Body.Close()
        respBytes, _ := io.ReadAll(resp.Body)

        if *debug {
            log.Println("===== API 3 Response Start =====")
            log.Printf("HTTP %d %s", resp.StatusCode, resp.Status)
            for k, vs := range resp.Header {
                for _, v := range vs {
                    log.Printf("  %s: %s", k, v)
                }
            }
            log.Printf("Body: %s", string(respBytes))
            log.Println("===== API 3 Response End =====")
        } else {
            log.Printf("[FORWARD] id=%s ohdCode=%s status=%d api3_http=%d cost=%v",
                cb.TaskID, ohdCode, status, resp.StatusCode, time.Since(start))
        }

        // Proxy API 3 response headers and body back to the original caller
        for k, vs := range resp.Header {
            for _, v := range vs {
                w.Header().Add(k, v)
            }
        }
        w.WriteHeader(resp.StatusCode)
        w.Write(respBytes)
    })

    log.Fatal(http.ListenAndServe(*listen, nil))
}

// Write API3 error-style response
func writeApi3Error(w http.ResponseWriter, code int, msg string) {
    res := map[string]interface{}{
        "result": false,
        "code":   code,
        "msg":    msg,
        "data":   nil,
        "object": nil,
    }
    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(200)
    bs, _ := json.Marshal(res)
    w.Write(bs)
}

// DES/ECB/PKCS5Padding + base64 encode
func desEncrypt(plain []byte, key string) (string, error) {
    k := []byte(key)
    if len(k) != 8 {
        return "", fmt.Errorf("DES key must be 8 bytes")
    }
    block, err := des.NewCipher(k)
    if err != nil {
        return "", err
    }
    padded := pkcs5Padding(plain, block.BlockSize())
    encrypted := make([]byte, len(padded))
    for bs, be := 0, block.BlockSize(); bs < len(padded); bs, be = bs+block.BlockSize(), be+block.BlockSize() {
        block.Encrypt(encrypted[bs:be], padded[bs:be])
    }
    return base64.StdEncoding.EncodeToString(encrypted), nil
}

func desDecrypt(b64data string, key string) (string, error) {
    k := []byte(key)
    if len(k) != 8 {
        return "", fmt.Errorf("DES key must be 8 bytes")
    }
    block, err := des.NewCipher(k)
    if err != nil {
        return "", err
    }
    decoded, err := base64.StdEncoding.DecodeString(b64data)
    if err != nil {
        return "", err
    }
    if len(decoded)%block.BlockSize() != 0 {
        return "", fmt.Errorf("input data is not DES block aligned")
    }
    decrypted := make([]byte, len(decoded))
    for bs, be := 0, block.BlockSize(); bs < len(decoded); bs, be = bs+block.BlockSize(), be+block.BlockSize() {
        block.Decrypt(decrypted[bs:be], decoded[bs:be])
    }
    unpadded := pkcs5Unpadding(decrypted)
    return string(unpadded), nil
}

func pkcs5Padding(src []byte, blockSize int) []byte {
    pad := blockSize - len(src)%blockSize
    padding := bytes.Repeat([]byte{byte(pad)}, pad)
    return append(src, padding...)
}

func pkcs5Unpadding(src []byte) []byte {
    l := len(src)
    if l == 0 {
        return src
    }
    pad := int(src[l-1])
    if pad < 1 || pad > l {
        return src // fallback
    }
    return src[:l-pad]
}


